﻿#pragma once
#include <map>
#include <memory>
#include <thread>
#include <atomic>
#include <chrono>
#include <iostream>
#include <boost/asio.hpp>
#include "types.hpp"
#include "modbus_slave.hpp"
#include "data_persistence.hpp"

// 协议类型转字符串
std::string protocolToString(ProtocolType protocol) {
    switch (protocol) {
    case ProtocolType::MODBUS: return "Modbus";
    case ProtocolType::MODBUS_RTU: return "Modbus RTU";
    case ProtocolType::MODBUS_TCP: return "Modbus TCP";
    default: return "Unknown";
    }
}

class ModbusServer {
public:
    ModbusServer(boost::asio::io_context& io_context,
        TransportType transport,
        ProtocolType protocol,
        uint16_t port = 502)
        : io_context_(io_context)
        , transport_type_(transport)
        , protocol_type_(protocol)
        , port_(port) {
    }

    ~ModbusServer() {
        stop();
    }

    void start() {
        if (running_) {
            std::cout << "服务器已在运行" << std::endl;
            return;
        }

        running_ = true;

        try {
            if (transport_type_ == TransportType::TCP) {
                startTCPServer();
            }
            else {
                startUDPServer();
            }

            // 启动IO上下文运行线程
            io_thread_ = std::thread(&ModbusServer::runIOContext, this);

            std::cout << "Modbus服务器已启动 ("
                << (transport_type_ == TransportType::TCP ? "TCP" : "UDP")
                << ", 端口: " << port_
                << ", 协议: " << protocolToString(protocol_type_)
                << ")" << std::endl;

        }
        catch (const std::exception& e) {
            std::cerr << "启动服务器失败: " << e.what() << std::endl;
            running_ = false;
            throw;
        }
    }

    void stop() {
        if (!running_) {
            return;
        }

        running_ = false;

        // 停止接受新连接
        if (tcp_acceptor_) {
            boost::system::error_code ec;
            tcp_acceptor_->cancel(ec);
            tcp_acceptor_->close(ec);
        }

        // 停止UDP socket
        if (udp_socket_) {
            boost::system::error_code ec;
            udp_socket_->cancel(ec);
            udp_socket_->close(ec);
        }

        // 停止IO上下文
        io_context_.stop();

        // 等待IO线程结束
        if (io_thread_.joinable()) {
            io_thread_.join();
        }

        std::cout << "Modbus服务器已停止" << std::endl;
    }

    void addSlave(uint8_t slave_id, const ModbusData& initial_data = ModbusData()) {
        auto it = slaves_.find(slave_id);
        if (it != slaves_.end()) {
            std::cout << "从站 " << static_cast<int>(slave_id) << " 已存在" << std::endl;
            return;
        }

        // 更新持久化数据
        SlaveConfig config;
        config.slave_id = slave_id;
        config.data = initial_data;

        // 如果数据为空，使用默认大小初始化
        if (config.data.coils.empty()) {
            config.data.coils.resize(ModbusData::DEFAULT_COILS, false);
        }
        if (config.data.discrete_inputs.empty()) {
            config.data.discrete_inputs.resize(ModbusData::DEFAULT_DISCRETE_INPUTS, false);
        }
        if (config.data.holding_registers.empty()) {
            config.data.holding_registers.resize(ModbusData::DEFAULT_HOLDING_REGISTERS, 0);
        }
        if (config.data.input_registers.empty()) {
            config.data.input_registers.resize(ModbusData::DEFAULT_INPUT_REGISTERS, 0);
        }

        persistence_data_.slaves[slave_id] = config;

        // 创建从站实例
        slaves_[slave_id] = std::make_unique<ModbusSlave>(slave_id, persistence_data_.slaves[slave_id].data);

        std::cout << "添加从站 " << static_cast<int>(slave_id) << std::endl;
    }

    void removeSlave(uint8_t slave_id) {
        auto it = slaves_.find(slave_id);
        if (it != slaves_.end()) {
            slaves_.erase(it);
            persistence_data_.slaves.erase(slave_id);
            std::cout << "移除从站 " << static_cast<int>(slave_id) << std::endl;
        }
    }

    bool saveData(const std::string& filename) {
        persistence_data_.timestamp = static_cast<uint32_t>(std::chrono::system_clock::now().time_since_epoch().count());
        return DataPersistence::saveToFile(filename, persistence_data_);
    }

    bool loadData(const std::string& filename) {
        PersistenceData new_data;
        if (!DataPersistence::loadFromFile(filename, new_data)) {
            return false;
        }

        // 清除现有从站
        slaves_.clear();
        persistence_data_ = new_data;

        // 重新创建从站实例
        for (auto& [slave_id, config] : persistence_data_.slaves) {
            slaves_[slave_id] = std::make_unique<ModbusSlave>(slave_id, config.data);
        }

        std::cout << "从文件加载数据成功，共 " << slaves_.size() << " 个从站" << std::endl;
        return true;
    }

private:
    void startTCPServer() {
        try {
            tcp_acceptor_ = std::make_unique<boost::asio::ip::tcp::acceptor>(
                io_context_, boost::asio::ip::tcp::endpoint(boost::asio::ip::tcp::v4(), port_));

            std::cout << "TCP服务器监听端口: " << port_ << std::endl;
            doTCPAccept();
        }
        catch (const std::exception& e) {
            std::cerr << "启动TCP服务器失败: " << e.what() << std::endl;
            throw;
        }
    }

    void startUDPServer() {
        try {
            udp_socket_ = std::make_unique<boost::asio::ip::udp::socket>(
                io_context_, boost::asio::ip::udp::endpoint(boost::asio::ip::udp::v4(), port_));

            std::cout << "UDP服务器监听端口: " << port_ << std::endl;
            doUDPRead();
        }
        catch (const std::exception& e) {
            std::cerr << "启动UDP服务器失败: " << e.what() << std::endl;
            throw;
        }
    }

    void doTCPAccept() {
        if (!running_) return;

        // 创建新的socket用于接受连接
        auto new_socket = std::make_shared<boost::asio::ip::tcp::socket>(io_context_);

        tcp_acceptor_->async_accept(*new_socket,
            [this, new_socket](boost::system::error_code ec) {
                if (!ec) {
                    std::cout << "TCP客户端连接: "
                        << new_socket->remote_endpoint().address().to_string()
                        << ":" << new_socket->remote_endpoint().port() << std::endl;

                    // 为每个连接启动独立的读取循环
                    doTCPRead(new_socket);
                }
                else {
                    if (ec != boost::asio::error::operation_aborted) {
                        std::cerr << "接受连接错误: " << ec.message() << std::endl;
                    }
                }

                // 继续接受新的连接
                if (running_) {
                    doTCPAccept();
                }
            });
    }

    void doTCPRead(std::shared_ptr<boost::asio::ip::tcp::socket> socket) {
        if (!running_) return;

        auto buffer = std::make_shared<std::vector<uint8_t>>(1024); // 增大缓冲区

        socket->async_read_some(boost::asio::buffer(*buffer),
            [this, socket, buffer](boost::system::error_code ec, std::size_t length) {
                if (!ec) {
                    buffer->resize(length);

                    // 调试信息
                    printf("TCP RX ");
                    for (size_t i = 0; i < buffer->size(); i++) {
                        printf(i == buffer->size() - 1 ? "%02X\n" : "%02X ", (*buffer)[i]);
                    }

                    auto send_callback = [socket](const std::shared_ptr<std::vector<uint8_t>> response) {
                        if (!response->empty()) {
                            boost::asio::async_write(*socket,
                                boost::asio::buffer(*response),
                                [socket, response](boost::system::error_code ec, std::size_t length) {
                                    if (ec) {
                                        std::cerr << "发送响应错误: " << ec.message() << std::endl;
                                    }
                                    else {
                                        // 调试信息
                                        printf("TCP TX ");
                                        for (size_t i = 0; i < length; i++) {
                                            printf(i == length - 1 ? "%02X\n" : "%02X ", (*response)[i]);
                                        }
                                    }
                                });
                        }
                        };

                    handleRequest(*buffer, send_callback);

                    // 继续读取
                    doTCPRead(socket);
                }
                else {
                    if (ec == boost::asio::error::eof) {
                        std::cout << "TCP客户端正常断开连接" << std::endl;
                    }
                    else if (ec == boost::asio::error::connection_reset) {
                        std::cout << "TCP连接被重置" << std::endl;
                    }
                    else if (ec != boost::asio::error::operation_aborted) {
                        std::cerr << "TCP读取错误: " << ec.message() << std::endl;
                    }
                    // 连接关闭，不需要继续读取
                }
            });
    }

    void doUDPRead() {
        if (!running_) return;

        auto buffer = std::make_shared<std::vector<uint8_t>>(1024);
        auto remote_endpoint = std::make_shared<boost::asio::ip::udp::endpoint>();

        udp_socket_->async_receive_from(
            boost::asio::buffer(*buffer),
            *remote_endpoint,
            [this, buffer, remote_endpoint](boost::system::error_code ec, std::size_t length) {
                if (!ec) {
                    buffer->resize(length);

                    // 调试信息
                    printf("UDP RX ");
                    for (size_t i = 0; i < buffer->size(); i++) {
                        printf(i == buffer->size() - 1 ? "%02X\n" : "%02X ", (*buffer)[i]);
                    }

                    auto send_callback = [this, remote_endpoint](const std::shared_ptr<std::vector<uint8_t>> response) {
                        if (!response->empty()) {
                            udp_socket_->async_send_to(
                                boost::asio::buffer(*response),
                                *remote_endpoint,
                                [remote_endpoint, response](boost::system::error_code ec, std::size_t length) {
                                    if (ec) {
                                        std::cerr << "发送UDP响应错误: " << ec.message()
                                            << " 到 " << remote_endpoint->address().to_string()
                                            << ":" << remote_endpoint->port() << std::endl;
                                    }
                                    else {
                                        // 调试信息
                                        printf("UDP TX ");
                                        for (size_t i = 0; i < length; i++) {
                                            printf(i == length - 1 ? "%02X\n" : "%02X ", (*response)[i]);
                                        }
                                    }
                                });
                        }
                        };

                    handleRequest(*buffer, send_callback);

                    // 重要：继续监听下一个UDP数据包
                    doUDPRead();

                }
                else {
                    if (ec != boost::asio::error::operation_aborted) {
                        std::cerr << "UDP接收错误: " << ec.message() << std::endl;

                        // 错误时也继续监听，但稍作延迟避免快速重试
                        if (running_) {
                            std::this_thread::sleep_for(std::chrono::milliseconds(100));
                            doUDPRead();
                        }
                    }
                }
            });
    }

    void handleRequest(const std::vector<uint8_t>& request,
        std::function<void(const std::shared_ptr<std::vector<uint8_t>>)> send_callback) {
        if (request.empty()) {
            return;
        }

        // 提取从站ID
        uint8_t slave_id = 0;
        if (protocol_type_ == ProtocolType::MODBUS_TCP && request.size() >= 7) {
            slave_id = request[6]; // TCP协议中单元标识符在第7个字节
        }
        else if (protocol_type_ == ProtocolType::MODBUS_RTU && request.size() >= 1) {
            slave_id = request[0]; // RTU协议中从站ID在第一个字节
        }
        else if (protocol_type_ == ProtocolType::MODBUS && request.size() >= 1) {
            slave_id = request[0];
        }

        // 查找对应的从站
        auto it = slaves_.find(slave_id);
        if (it != slaves_.end()) {
            std::shared_ptr<std::vector<uint8_t>> response =std::make_shared<std::vector<uint8_t>>(
                it->second->processRequest(request, protocol_type_));
            if (!response->empty()) {
                send_callback(response);
            }
        }
    }

    void runIOContext() {
        while (running_) {
            try {
                // 重置io_context，因为run()在停止后需要重置才能再次运行
                io_context_.restart();
                io_context_.run();
                break; // io_context.run() 正常完成
            }
            catch (const std::exception& e) {
                std::cerr << "IO上下文运行错误: " << e.what() << std::endl;

                // 短暂延迟后重试
                if (running_) {
                    std::this_thread::sleep_for(std::chrono::seconds(1));
                }
            }
        }

        std::cout << "IO上下文运行线程结束" << std::endl;
    }

private:
    boost::asio::io_context& io_context_;
    TransportType transport_type_;
    ProtocolType protocol_type_;
    uint16_t port_;

    std::unique_ptr<boost::asio::ip::tcp::acceptor> tcp_acceptor_;
    std::unique_ptr<boost::asio::ip::udp::socket> udp_socket_;

    std::map<uint8_t, std::unique_ptr<ModbusSlave>> slaves_;
    PersistenceData persistence_data_;

    std::atomic<bool> running_{ false };
    std::thread io_thread_;
};
